package com.citro.agents;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.connector.jdbc.JdbcConnectionOptions;
import org.apache.flink.connector.jdbc.JdbcExecutionOptions;
import org.apache.flink.connector.jdbc.JdbcSink;
import org.apache.flink.connector.jdbc.JdbcStatementBuilder;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.connectors.kinesis.FlinkKinesisConsumer;
import org.apache.flink.streaming.connectors.kinesis.config.ConsumerConfigConstants;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.Properties;

/**
 * Flink Job for Behavior Analysis using LLM
 *
 * Architecture:
 * Kinesis Stream -> Parse JSON -> LLM Analysis -> PostgreSQL Sink
 */
public class BehaviorAnalysisJob {

    private static final Logger LOG = LoggerFactory.getLogger(BehaviorAnalysisJob.class);
    private static final ObjectMapper objectMapper = new ObjectMapper();

    // Configuration from environment variables
    private static final String KINESIS_STREAM_NAME = getEnv("KINESIS_STREAM_NAME", "click-events-stream");
    private static final String KINESIS_REGION = getEnv("AWS_REGION", "us-east-1");
    private static final String KINESIS_ENDPOINT = getEnv("KINESIS_ENDPOINT", "http://localstack:4566");

    private static final String POSTGRES_HOST = getEnv("POSTGRES_HOST", "postgres");
    private static final String POSTGRES_PORT = getEnv("POSTGRES_PORT", "5432");
    private static final String POSTGRES_DB = getEnv("POSTGRES_DB", "browse_data");
    private static final String POSTGRES_USER = getEnv("POSTGRES_USER", "admin");
    private static final String POSTGRES_PASSWORD = getEnv("POSTGRES_PASSWORD", "admin123");

    private static final String LLM_PROVIDER = getEnv("LLM_PROVIDER", "deepseek");
    private static final String LLM_API_KEY = getEnv("LLM_API_KEY", "");
    private static final String LLM_API_URL = getEnv("LLM_API_URL", "https://api.deepseek.com/v1/chat/completions");
    private static final String LLM_MODEL = getEnv("LLM_MODEL", "deepseek-chat");
    private static final int LLM_TIMEOUT = Integer.parseInt(getEnv("LLM_TIMEOUT", "30"));

    // Simple prompt template (can be loaded from file later)
    private static final String PROMPT_TEMPLATE =
        "Analyze this browser event and describe the user's behavior:\n\n" +
        "Event Type: {action}\n" +
        "URL: {url}\n" +
        "Page Title: {title}\n" +
        "Semantic Structure: {semantic_tree}\n\n" +
        "Respond in JSON format:\n" +
        "{\"behavior_type\": \"search|navigation|interaction|content_creation|social_action\", " +
        "\"behavior_description\": \"brief description\", " +
        "\"confidence\": 0.9}";

    public static void main(String[] args) throws Exception {
        LOG.info("Starting Flink Agents Behavior Analysis Job");
        LOG.info("Kinesis Stream: {}", KINESIS_STREAM_NAME);
        LOG.info("LLM Provider: {}", LLM_PROVIDER);

        // Set up the streaming execution environment
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        // Enable checkpointing for fault tolerance
        env.enableCheckpointing(60000); // checkpoint every 60 seconds

        // Configure Kinesis Consumer
        Properties kinesisConsumerConfig = new Properties();
        kinesisConsumerConfig.setProperty(ConsumerConfigConstants.AWS_REGION, KINESIS_REGION);
        kinesisConsumerConfig.setProperty(ConsumerConfigConstants.AWS_ENDPOINT, KINESIS_ENDPOINT);
        kinesisConsumerConfig.setProperty(ConsumerConfigConstants.STREAM_INITIAL_POSITION, "LATEST");

        // Use dummy credentials for LocalStack
        kinesisConsumerConfig.setProperty(ConsumerConfigConstants.AWS_CREDENTIALS_PROVIDER, "BASIC");
        kinesisConsumerConfig.setProperty("aws.credentials.provider.basic.accesskeyid", "test");
        kinesisConsumerConfig.setProperty("aws.credentials.provider.basic.secretkey", "test");

        // Create Kinesis source
        FlinkKinesisConsumer<String> kinesisSource = new FlinkKinesisConsumer<>(
                KINESIS_STREAM_NAME,
                new SimpleStringSchema(),
                kinesisConsumerConfig
        );

        // Read from Kinesis
        DataStream<String> rawEvents = env.addSource(kinesisSource)
                .name("Kinesis Source");

        // Parse JSON to BehaviorEvent
        DataStream<BehaviorEvent> parsedEvents = rawEvents
                .map(new JsonParseFunction())
                .name("Parse JSON");

        // Analyze with LLM
        DataStream<BehaviorAnalysisResult> analyzedEvents = parsedEvents
                .map(new LLMAnalysisFunction(LLM_PROVIDER, LLM_API_KEY, LLM_API_URL, LLM_MODEL, LLM_TIMEOUT))
                .name("LLM Analysis");

        // Sink to PostgreSQL
        String jdbcUrl = String.format("jdbc:postgresql://%s:%s/%s",
                POSTGRES_HOST, POSTGRES_PORT, POSTGRES_DB);

        analyzedEvents.addSink(JdbcSink.sink(
                "INSERT INTO user_behaviors " +
                "(user_id, session_id, behavior_type, behavior_description, url, confidence, analysis_method) " +
                "VALUES (?, ?, ?, ?, ?, ?, ?)",
                new JdbcStatementBuilder<BehaviorAnalysisResult>() {
                    @Override
                    public void accept(PreparedStatement ps, BehaviorAnalysisResult result) throws SQLException {
                        ps.setString(1, result.getUserId());
                        ps.setString(2, result.getSessionId());
                        ps.setString(3, result.getBehaviorType());
                        ps.setString(4, result.getBehaviorDescription());
                        ps.setString(5, result.getUrl());
                        ps.setDouble(6, result.getConfidence());
                        ps.setString(7, result.getAnalysisMethod());
                    }
                },
                JdbcExecutionOptions.builder()
                        .withBatchSize(10)
                        .withBatchIntervalMs(1000)
                        .withMaxRetries(3)
                        .build(),
                new JdbcConnectionOptions.JdbcConnectionOptionsBuilder()
                        .withUrl(jdbcUrl)
                        .withDriverName("org.postgresql.Driver")
                        .withUsername(POSTGRES_USER)
                        .withPassword(POSTGRES_PASSWORD)
                        .build()
        )).name("PostgreSQL Sink");

        // Execute the job
        env.execute("Behavior Analysis with LLM (Java)");
    }

    /**
     * Function to parse JSON string to BehaviorEvent
     */
    private static class JsonParseFunction implements MapFunction<String, BehaviorEvent> {
        @Override
        public BehaviorEvent map(String json) throws Exception {
            try {
                return objectMapper.readValue(json, BehaviorEvent.class);
            } catch (Exception e) {
                LOG.error("Failed to parse JSON: {}", json, e);
                // Return empty event to avoid job failure
                return new BehaviorEvent();
            }
        }
    }

    /**
     * Function to analyze behavior using LLM
     */
    private static class LLMAnalysisFunction implements MapFunction<BehaviorEvent, BehaviorAnalysisResult> {
        private transient LLMClient llmClient;
        private final String provider;
        private final String apiKey;
        private final String apiUrl;
        private final String model;
        private final int timeout;

        public LLMAnalysisFunction(String provider, String apiKey, String apiUrl, String model, int timeout) {
            this.provider = provider;
            this.apiKey = apiKey;
            this.apiUrl = apiUrl;
            this.model = model;
            this.timeout = timeout;
        }

        @Override
        public BehaviorAnalysisResult map(BehaviorEvent event) throws Exception {
            // Lazy initialization (per task)
            if (llmClient == null) {
                llmClient = new LLMClient(provider, apiKey, apiUrl, model, timeout);
            }

            return llmClient.analyzeEvent(event, PROMPT_TEMPLATE);
        }
    }

    /**
     * Utility to get environment variable with default
     */
    private static String getEnv(String key, String defaultValue) {
        String value = System.getenv(key);
        return (value != null && !value.isEmpty()) ? value : defaultValue;
    }
}
